{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric_temporal.nn.recurrent import A3TGCN2\n",
    "from torch_geometric_temporal.signal import temporal_signal_split\n",
    "import pandas as pd\n",
    "import json\n",
    "\n",
    "# GPU support\n",
    "DEVICE = torch.device('cpu') # cuda\n",
    "shuffle=True\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = pd.read_csv(r'C:\\Users\\edi\\GitHub\\pinns_opf\\data\\interim\\lines_5.csv')\n",
    "nodes = pd.read_csv(r'C:\\Users\\edi\\GitHub\\pinns_opf\\data\\interim\\nodes_5.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons_id = 0\n",
    "cons_id_node = {}\n",
    "\n",
    "pv_id = 0\n",
    "pv_id_node = {}\n",
    "\n",
    "ev_id = 0\n",
    "ev_id_node = {}\n",
    "\n",
    "dg_id = 0\n",
    "dg_id_node = {}\n",
    "\n",
    "ess_id = 0\n",
    "ess_id_node = {}\n",
    "for row in nodes.iterrows():\n",
    "\n",
    "    for j in range(len(row[1])):\n",
    "\n",
    "        if row[1][j] == 'cons':\n",
    "            cons_id += 1\n",
    "            cons_id_node[cons_id] = row[1][0]\n",
    "        \n",
    "        if row[1][j] == 'pv':\n",
    "            pv_id += 1\n",
    "            pv_id_node[pv_id] = row[1][0]\n",
    "        \n",
    "        if row[1][j] == 'ev':\n",
    "            ev_id += 1\n",
    "            ev_id_node[ev_id] = row[1][0]\n",
    "\n",
    "        if row[1][j] == 'ess':\n",
    "            ess_id += 1\n",
    "            ess_id_node[ess_id] = row[1][0]\n",
    "        \n",
    "        if row[1][j] == 'dg':\n",
    "            dg_id += 1\n",
    "            dg_id_node[dg_id] = row[1][0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('variable_data.json') as f:\n",
    "    dataset = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "for key in dataset.keys():\n",
    "\n",
    "    for node_id in dataset[key].keys():\n",
    "\n",
    "        if 'V' in key:\n",
    "            print(node_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "hetero_graph_dataset = dict()\n",
    "node_id = nodes['Nodes'].values\n",
    "for node in node_id:\n",
    "    hetero_graph_dataset[node] = []\n",
    "    \n",
    "for key in dataset.keys():\n",
    "\n",
    "    for node_id in dataset[key].keys():\n",
    "        \n",
    "        if 'cons' in key:\n",
    "            node_id_g = cons_id_node[int(node_id)]\n",
    "            hetero_graph_dataset[node_id_g].append(dataset[key][str(node_id_g)])\n",
    "        elif 'ev' in key:\n",
    "            node_id_g = ev_id_node[int(node_id)]\n",
    "            hetero_graph_dataset[node_id_g].append(dataset[key][str(node_id_g)])\n",
    "        elif 'pv' in key:\n",
    "            node_id_g = pv_id_node[int(node_id)]\n",
    "            hetero_graph_dataset[node_id_g].append(dataset[key][str(node_id)])\n",
    "        elif 'ess' in key:\n",
    "            node_id_g = ess_id_node[int(node_id)]\n",
    "            hetero_graph_dataset[node_id_g].append(dataset[key][str(node_id)])\n",
    "        elif 'dg' in key:\n",
    "            node_id_g = dg_id_node[int(node_id)]\n",
    "            hetero_graph_dataset[node_id_g].append(dataset[key][str(node_id)])\n",
    "        elif 'V' in key:\n",
    "            hetero_graph_dataset[int(node_id)].append(dataset[key][str(node_id)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 (9, 24)\n",
      "2 (9, 24)\n",
      "3 (9, 24)\n",
      "4 (9, 24)\n",
      "5 (9, 24)\n"
     ]
    }
   ],
   "source": [
    "list_ = []\n",
    "for node,value in hetero_graph_dataset.items():\n",
    "    hetero_graph_dataset[node] = np.array(value)\n",
    "    list_.append(np.array(value))\n",
    "    print(node,hetero_graph_dataset[node].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stack the arrays from the dictionary into the result array\n",
    "# Create an empty numpy array of size 24x5x9\n",
    "result_array = np.empty((24, 5, 9))\n",
    "\n",
    "for i, key in enumerate(hetero_graph_dataset.keys()):\n",
    "    result_array[:, i, :] = hetero_graph_dataset[key].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reshape the array to (120, 9) to apply normalization to each feature\n",
    "data_reshaped = np.reshape(result_array, (result_array.shape[0]*result_array.shape[1], result_array.shape[2]))\n",
    "\n",
    "# Calculate the mean and standard deviation along the 0th axis (rows)\n",
    "mean = np.mean(data_reshaped, axis=0)\n",
    "std = np.std(data_reshaped, axis=0)\n",
    "\n",
    "# Normalize the data using mean and standard deviation\n",
    "normalized_data = (data_reshaped - mean) / std\n",
    "\n",
    "# Reshape the normalized data back to the original shape\n",
    "normalized_data = np.reshape(normalized_data, (result_array.shape[0], result_array.shape[1], result_array.shape[2]))\n",
    "normalized_data[:,:,-1] = result_array[:,:,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "\n",
    "# # Assuming your array is called 'data'\n",
    "# # Shape: (24, 5, 9)\n",
    "\n",
    "# # Reshape the array to (120, 9) to apply normalization to each feature\n",
    "# data_reshaped = np.reshape(result_array, (result_array.shape[0]*result_array.shape[1], result_array.shape[2]))\n",
    "\n",
    "# # Calculate the minimum and maximum values along the 0th axis (rows)\n",
    "# min_vals = np.min(data_reshaped, axis=0)\n",
    "# max_vals = np.max(data_reshaped, axis=0)\n",
    "\n",
    "# # Normalize the data using min-max normalization\n",
    "# normalized_data = (data_reshaped - min_vals) / (max_vals - min_vals)\n",
    "\n",
    "# # Reshape the normalized data back to the original shape\n",
    "# normalized_data = np.reshape(normalized_data, (result_array.shape[0], result_array.shape[1], result_array.shape[2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_index = np.array([lines['From'].values,lines['To'].values])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_weigth = np.ones(edge_index.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = normalized_data[:-1,:,:]\n",
    "targets = normalized_data[1:,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "features_seq = []\n",
    "targets_seq = []\n",
    "\n",
    "for i in range(1000*len(features[:,1,1])):\n",
    "    features_seq.append(features[i%23,:,:])\n",
    "    targets_seq.append(targets[i%23,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric_temporal.signal import StaticGraphTemporalSignal\n",
    "\n",
    "dataset = StaticGraphTemporalSignal(\n",
    "    edge_index=edge_index, edge_weight=edge_weigth,\n",
    "    features=features_seq, targets=targets_seq\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1., 1., 1., 1., 1., 1.])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.__dict__['edge_weight']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric_temporal.nn.recurrent import DCRNN,GCLSTM\n",
    "\n",
    "class RecurrentGCN(torch.nn.Module):\n",
    "    def __init__(self, node_features):\n",
    "        super(RecurrentGCN, self).__init__()\n",
    "        self.recurrent = DCRNN(node_features, 32, 1)\n",
    "        self.linear = torch.nn.Linear(32, 9)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight):\n",
    "        h = self.recurrent(x, edge_index, edge_weight)\n",
    "        h = F.relu(h)\n",
    "        h = self.linear(h)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric_temporal.nn.recurrent import GConvLSTM\n",
    "\n",
    "class RecurrentGCN(torch.nn.Module):\n",
    "    def __init__(self, node_features):\n",
    "        super(RecurrentGCN, self).__init__()\n",
    "        self.recurrent = GConvLSTM(node_features, 32, 1, normalization='rw')\n",
    "        self.linear = torch.nn.Linear(32, 9)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight, h, c):\n",
    "        h_0, c_0 = self.recurrent(x, edge_index, edge_weight, h, c)\n",
    "        h = F.relu(h_0)\n",
    "        h = self.linear(h)\n",
    "        return h, h_0, c_0\n",
    "        \n",
    "model = RecurrentGCN(node_features=9)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "model = RecurrentGCN(node_features=9)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inequality_loss(v_pred, v_lower, v_upper):\n",
    "    \"\"\"\n",
    "    Computes the loss for inequality constraints v_lower <= v_pred <= v_upper.\n",
    "    \n",
    "    Arguments:\n",
    "    v_pred -- predicted vector of size 5x1\n",
    "    v_lower -- lower bound vector of size 5x1\n",
    "    v_upper -- upper bound vector of size 5x1\n",
    "    \n",
    "    Returns:\n",
    "    loss -- the computed loss\n",
    "    \"\"\"\n",
    "    # Compute violation of lower bound\n",
    "    violation_lower = torch.relu(v_lower - v_pred)\n",
    "    \n",
    "    # Compute violation of upper bound\n",
    "    violation_upper = torch.relu(v_pred - v_upper)\n",
    "    \n",
    "    # Compute the total violation\n",
    "    total_violation = torch.sum(violation_lower) + torch.sum(violation_upper)\n",
    "    \n",
    "    # Compute the hinge loss\n",
    "    hinge_loss = F.relu(total_violation)\n",
    "    \n",
    "    return hinge_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.9946, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 1/10 [03:59<35:52, 239.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.9176, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 2/10 [06:00<22:37, 169.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8440, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 3/10 [08:11<17:45, 152.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.7694, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 4/10 [10:07<13:47, 137.93s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6947, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 5/10 [12:06<10:55, 131.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6199, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 6/10 [14:01<08:21, 125.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6633, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 7/10 [16:33<06:42, 134.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5297, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 8/10 [19:32<04:57, 148.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5918, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 9/10 [21:22<02:16, 136.35s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.4733, grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [24:57<00:00, 149.76s/it]\n"
     ]
    }
   ],
   "source": [
    "model = RecurrentGCN(node_features=9)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "model.train()\n",
    "n = 5  # number of nodes\n",
    "v_lower = torch.zeros((n, 1))\n",
    "v_upper = torch.ones((n, 1))\n",
    "\n",
    "_lambda = 0.7\n",
    "\n",
    "for epoch in tqdm(range(10)):\n",
    "    cost = 0\n",
    "    h, c = None, None\n",
    "    for time, snapshot in enumerate(train_dataset):\n",
    "        y_hat, h, c = model(snapshot.x, snapshot.edge_index, snapshot.edge_attr, h, c)\n",
    "        v_pred = y_hat[:,-1]\n",
    "        cost = cost + torch.mean((y_hat-snapshot.y)**2) + _lambda*inequality_loss(v_pred, v_lower, v_upper)   \n",
    "    cost = cost / (time+1)\n",
    "    print(cost)\n",
    "    cost.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "opf_pinns_venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
